#!/usr/bin/env python3

import sys, getopt, os, tarfile, zlib, hashlib, json, shutil
sys.path.append(os.getcwd())
from data.firmware import ver

def usage():
    print('pack_firmware.py -i <inputpath> -o <outputfile>')
    
def main(argv):
    inputpath = ''
    outputfile = ''
    try:
        opts, args = getopt.getopt(argv,"hi:o:")
    except getopt.GetoptError:
        usage()
        sys.exit(2)
        
    for opt, arg in opts:
        if opt == '-h':
            usage()
            sys.exit()
        elif opt in ('-i'):
            inputpath = arg
        elif opt in ('-o'):
            outputfile = arg
            
    print("inputpath = ", inputpath)
    print("outputfile = ", outputfile)
    
    isPathExist = os.path.exists(inputpath)
    if (not isPathExist):
        print(inputpath + "isn't exist.")
        sys.exit()
    
    outputfileTar = outputfile + ".tar"
    isOutputExist = os.path.isfile(outputfileTar)
    print(outputfileTar , isOutputExist)
    if (isOutputExist):
        os.remove(outputfileTar)

    # mpy-cross 预处理
    if (os.path.exists('./Release/firmware/')):
        shutil.rmtree('./Release/firmware/')
    os.mkdir('./Release/firmware')
    for entry in os.listdir(inputpath):
        fn = os.path.join(inputpath, entry)
        if (entry.endswith('.py')):
            if (entry == 'main.py'):
                shutil.copyfile(fn, './Release/firmware/main.py')
            else:
                print('./tools/mpy-cross ' + fn + ' -o ./Release/firmware/' + entry[:-3] + '.mpy')
                os.system('./tools/mpy-cross ' + fn + ' -o ./Release/firmware/' + entry[:-3] + '.mpy')

    inputpath = './Release/firmware'
    
    # TODO：计算目录下所有文件的MD5，合并到一个文件中
    filelistMD5 = open(os.path.join(inputpath, 'filelist.md5'), 'w+')
    
    filelistMD5Payload = {}
    for entry in os.listdir(inputpath):
        fn = os.path.join(inputpath, entry)
        if (os.path.isfile(fn) and (entry != 'filelist.md5')):
            # MD5Sum
            filelistMD5Payload[entry] = hashlib.md5(open(fn, 'rb').read()).hexdigest()
    
    filelistMD5.write(json.dumps(filelistMD5Payload))
    filelistMD5.flush()
    filelistMD5.close()
            
    
    # 先tar目录
    tarOutput = tarfile.open(outputfileTar, 'x')
        
    # TODO：检查，这个目录下的根目录里，必须有个main.py文件
    isMainExist = os.path.isfile(inputpath + "/main.py")
    if (not isMainExist):
        print("Error: Need main.py in input path.")
        sys.exit()

    tarOutput.add(inputpath, '')
    tarOutput.close()
    
    # 用zlib特定方式压缩
    tarf = open(outputfileTar, 'rb')
    tarPayload = tarf.read()
    zc = zlib.compressobj(wbits=25)
    zPayload = zc.compress(tarPayload)
    zPayload += zc.flush()
    
    outputfileBin = outputfile + ".bin"
    isOutputExist = os.path.isfile(outputfileBin)
    print(outputfileBin , isOutputExist)
    if (isOutputExist):
        os.remove(outputfileBin)
        
    cFile = open(outputfile + ".bin", 'wb')
    cFile.write(zPayload)
    cFile.flush()
    cFile.close()
    
    # 删除临时的tar文件
    
    # 计算MD5
    firmwareMD5 = open(outputfileBin + ".md5", 'w+')
    firmwareMD5Payload = hashlib.md5(zPayload).hexdigest()
    firmwareMD5.write(firmwareMD5Payload)
    firmwareMD5.write('\n')
    firmwareMD5.write("{0:0{1}x}".format(ver.ver,8))
    firmwareMD5.flush()
    firmwareMD5.close()
    
if __name__ == "__main__":
    main(sys.argv[1:])
    